import numpy as np
import albumentations as A

import plotly.graph_objects as go
from plotly.subplots import make_subplots

def read_obj(filename):
    triangles = []
    vertices = []
    with open(filename) as file:
        for line in file:
            components = line.strip(' \n').split(' ')
            if components[0] == "f":  # face data
                # e.g. "f 1/1/1/ 2/2/2 3/3/3 4/4/4 ..."
                indices = list(
                    map(lambda c: int(c.split('/')[0]) - 1, components[1:]))
                for i in range(0, len(indices) - 2):
                    triangles.append(indices[i:i + 3])
            elif components[0] == "v":  # vertex data
                # e.g. "v  30.2180 89.5757 -76.8089"
                vertex = list(map(lambda c: float(c), components[1:]))
                vertices.append(vertex)
    return np.array(vertices), np.array(triangles)

def show_augmentation(original, augmented):
    fig = make_subplots(rows=1,
                        cols=2,
                        specs=[[{
                            'type': 'scatter3d'
                        }, {
                            'type': 'scatter3d'
                        }]])

    before = go.Scatter3d(x=original[:, 0],
                          y=original[:, 1],
                          z=original[:, 2],
                          mode="markers",
                          marker=dict(size=2,
                                      color=original[:, 3]))

    after = go.Scatter3d(x=augmented[:, 0],
                         y=augmented[:, 1],
                         z=augmented[:, 2],
                         mode="markers",
                         marker=dict(size=2,
                                     color=augmented[:, 3]))
    fig.add_trace(before, row=1, col=1)
    fig.add_trace(after, row=1, col=2)
    fig.show()

teapot, _ = read_obj("images/teapot.obj")
teapot = np.hstack((teapot, np.zeros((len(teapot), 1))))
teapot[:,[0, 1, 2]] = teapot[:,[0, 2, 1]]
teapot[teapot[:, 0] > 2, 3] = 1
teapot[teapot[:, 0] < -2, 3] = 2
teapot[teapot[:, 2] > 2.5, 3] = 3

augmentation = A.Scale3d(scale_limit=[0.1, 0.1, 0.1], bias=[1.5, 1, 1], always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.RotateAroundAxis3d(axis=[0, 1, 0], rotation_limit=np.pi/3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.Crop3d(x_min=0, x_max=3, y_min=0, y_max=2, z_min=0, z_max=3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.RandomMove3d(x_min=5, x_max=7, y_min=10, y_max=20, z_min=0, z_max=3, always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.Center3d(always_apply=True)
augmented_teapot = augmentation(points=augmented_teapot)["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.Flip3d(axis=[0,0,1], always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)

augmentation = A.RandomDropout3d(always_apply=True)
augmented_teapot = augmentation(points=teapot.copy())["points"]
show_augmentation(teapot, augmented_teapot)